/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022.
 * Description: support carry dim kernel measurements across kexec
 * Author: chenjingwen6 <chenjingwen6@huawei.com>
 * Create: 2022-09-30
*/
#include <linux/mutex.h>
#include <linux/kernel.h>
#include <linux/vmalloc.h>
#include <linux/memblock.h>
#include <linux/libfdt.h>
#include <linux/of.h>
#include <linux/rtos/dim_kernel.h>

#define DIM_KEXEC_FDT_PROP_NAME "linux,dim-kexec-buffer"
#define DIM_DEBUG_DUMP_SIZE 100
#define DIM_DEBUG_DUMP_ROWSIZE 16
#define DIM_DEBUG_DUMP_GROUPSIZE 1

DEFINE_MUTEX(dim_kexec_mutex);
static void *dim_kexec_buffer;
static size_t dim_kexec_buffer_size;
static fn_dim_kexec_store_cb dim_kexec_store_cb;

/*
 * set kexec store callback
 * When kexec load, dim_add_kexec_buffer will call store_callback
 * to store dim measurements to buffer and buffer can be carried
 * across kexec
 */
int dim_set_kexec_store_cb(fn_dim_kexec_store_cb store_callback)
{
	int r = 0;

	mutex_lock(&dim_kexec_mutex);
	if (dim_kexec_store_cb) {
		r = -EEXIST;
		pr_err("DIM: kexec store callback is set already, please clear it before set\n");
		goto out_unlock;
	}

	if (!store_callback || !func_ptr_is_kernel_text(store_callback)) {
		r = -EINVAL;
		pr_err("DIM: %s Invalid store_callback\n", __func__);
		goto out_unlock;
	}

	dim_kexec_store_cb = store_callback;

out_unlock:
	mutex_unlock(&dim_kexec_mutex);
	return r;
}
EXPORT_SYMBOL_NS(dim_set_kexec_store_cb, HW_RTOS_NS);

void dim_clear_kexec_store_cb(void)
{
	mutex_lock(&dim_kexec_mutex);
	dim_kexec_store_cb = NULL;
	mutex_unlock(&dim_kexec_mutex);
}
EXPORT_SYMBOL_NS(dim_clear_kexec_store_cb, HW_RTOS_NS);

void dim_restore_from_kexec(fn_dim_kexec_restore_cb restore_callback)
{
	mutex_lock(&dim_kexec_mutex);
	if (!restore_callback || !func_ptr_is_kernel_text(restore_callback))
		goto out_unlock;

	/*
	 * dim kexec buffer is empty, because:
	 * 1. the kernel is not booted from kexec
	 * 2. no kexec buffer carry across kexec
	 * 3. there was one already restore measurements
	 */
	if (!dim_kexec_buffer || dim_kexec_buffer_size == 0) {
		pr_info("DIM: dim_kexec_buffer is empty, no need to restore measurement from kexec\n");
		goto out_unlock;
	}

	restore_callback(dim_kexec_buffer, dim_kexec_buffer_size);
	/* free buffer after someone already got it */
	vfree(dim_kexec_buffer);
	dim_kexec_buffer = NULL;
	dim_kexec_buffer_size = 0;
out_unlock:
	mutex_unlock(&dim_kexec_mutex);
}
EXPORT_SYMBOL_NS(dim_restore_from_kexec, HW_RTOS_NS);

void dim_add_kexec_buffer(struct kimage *image)
{
	int r;
	struct kexec_buf kbuf = { .image = image, .buf_align = PAGE_SIZE,
				  .buf_min = 0, .buf_max = ULONG_MAX,
				  .top_down = true };

	mutex_lock(&dim_kexec_mutex);
	if (!dim_kexec_store_cb || !func_ptr_is_kernel_text(dim_kexec_store_cb)) {
		pr_info("DIM: no register for kexec, no need to add dim to kexec buffer\n");
		goto out_unlock;
	}

	r = dim_kexec_store_cb(&kbuf);
	if (r < 0) {
		pr_err("DIM: dim_kexec_store_cb fail, cannot carry dim measurement across kexec, r:%d\n", r);
		goto out_unlock;
	}

	r = kexec_add_buffer(&kbuf);
	if (r < 0) {
		vfree(kbuf.buffer);
		kbuf.buffer = NULL;
		pr_err("DIM: Error passing over kexec_add_buffer, r:%d\n", r);
		goto out_unlock;
	}

	image->dim_buffer = kbuf.buffer;     /* virtual buffer of DIM measurement */
	image->dim_buffer_addr = kbuf.mem;   /* phys addr across kexec */
	image->dim_buffer_size = kbuf.memsz;

	pr_info("DIM: kexec buffer memsz:%zu\n", kbuf.memsz);
	print_hex_dump(KERN_INFO, "DIM dump: ", DUMP_PREFIX_NONE, DIM_DEBUG_DUMP_ROWSIZE, DIM_DEBUG_DUMP_GROUPSIZE,
				   kbuf.buffer, kbuf.bufsz < DIM_DEBUG_DUMP_SIZE ? kbuf.bufsz : DIM_DEBUG_DUMP_SIZE,
				   true);

out_unlock:
	mutex_unlock(&dim_kexec_mutex);
}

/*
 * setup_dim_buffer - add dim buffer information to the fdt
 * @image:		kexec image being loaded.
 * @fdt:		Flattened device tree for the next kernel.
 * @chosen_node:	Offset to the chosen node.
 *
 * Return: 0 on success, or negative errno on error.
 */
int setup_dim_buffer(const struct kimage *image, void *fdt, int chosen_node)
{
	int ret;

	if (image->dim_buffer_size == 0 || image->dim_buffer_addr == 0)
		return 0;

	ret = fdt_appendprop_addrrange(fdt, 0, chosen_node,
				       DIM_KEXEC_FDT_PROP_NAME,
				       image->dim_buffer_addr,
				       image->dim_buffer_size);
	if (ret < 0)
		return ret;

	ret = fdt_add_mem_rsv(fdt, image->dim_buffer_addr, image->dim_buffer_size);
	if (ret < 0)
		return ret;

	pr_info("DIM: setup dim buffer size in DTB: 0x%zx\n", image->dim_buffer_size);
	return 0;
}

/*
 * remove_dim_buffer - remove the DIM buffer property and reservation from @fdt
 *
 * @fdt: Flattened Device Tree to update
 * @chosen_node: Offset to the chosen node in the device tree
 *
 * The DIM measurement buffer is of no use to a subsequent kernel, so we always
 * remove it from the device tree.
 */
void remove_dim_buffer(void *fdt, int chosen_node)
{
	int ret, len;
	unsigned long addr;
	size_t size;
	const void *prop;

	prop = fdt_getprop(fdt, chosen_node, DIM_KEXEC_FDT_PROP_NAME, &len);
	if (!prop)
		return;

	ret = dim_helper_do_get_kexec_buffer(prop, len, &addr, &size);
	fdt_delprop(fdt, chosen_node, DIM_KEXEC_FDT_PROP_NAME);
	if (ret)
		return;

	ret = dim_helper_fdt_find_and_del_mem_rsv(fdt, addr, size);
	if (ret < 0)
		pr_info("DIM: Removed old DIM buffer reservation.\n");
}

/*
 * dim_get_kexec_buffer - get DIM buffer from the previous kernel
 * @addr: On successful return, set to point to the buffer contents
 * @size: On successful return, set to the buffer size
 *
 * Return: 0 on success, negative errno on error
 */
int __init dim_get_kexec_buffer(void **addr, size_t *size)
{
	int ret, len;
	unsigned long tmp_addr;
	unsigned long start_pfn, end_pfn;
	size_t tmp_size;
	const void *prop;

	prop = of_get_property(of_chosen, DIM_KEXEC_FDT_PROP_NAME, &len);
	if (!prop)
		return -ENOENT;

	ret = dim_helper_do_get_kexec_buffer(prop, len, &tmp_addr, &tmp_size);
	if (ret < 0)
		return ret;

	/* Do some sanity on the returned size for the dim-kexec buffer */
	if (tmp_size == 0)
		return -ENOENT;

	/*
	 * Calculate the PFNs for the buffer and ensure
	 * they are with in addressable memory.
	 */
	start_pfn = PHYS_PFN(tmp_addr);
	end_pfn = PHYS_PFN(tmp_addr + tmp_size - 1);
	if (!page_is_ram(start_pfn) || !page_is_ram(end_pfn)) {
		pr_warn("DIM: DIM buffer at 0x%lx, size = 0x%zx beyond memory\n",
			tmp_addr, tmp_size);
		return -EINVAL;
	}

	*addr = __va(tmp_addr);
	*size = tmp_size;
	pr_info("DIM: %s, dim buffer size: 0x%zx\n", __func__, tmp_size);
	return 0;
}

/*
 * dim_free_kexec_buffer - free memory used by the DIM buffer
 */
int __init dim_free_kexec_buffer(void)
{
	int ret;
	unsigned long addr;
	size_t size;
	struct property *prop;

	prop = of_find_property(of_chosen, DIM_KEXEC_FDT_PROP_NAME, NULL);
	if (!prop)
		return -ENOENT;

	ret = dim_helper_do_get_kexec_buffer(prop->value, prop->length, &addr, &size);
	if (ret < 0)
		return ret;

	ret = of_remove_property(of_chosen, prop);
	if (ret < 0)
		return ret;

	return memblock_free(addr, size);
}

/*
 * Restore the measurement list from the previous kernel.
 */
static void __init dim_load_kexec_buffer(void)
{
	void *kexec_buffer = NULL;
	size_t kexec_buffer_size = 0;
	int r;

	r = dim_get_kexec_buffer(&kexec_buffer, &kexec_buffer_size);
	switch (r) {
	case 0:
		/* temporarily kexec buffer store in vmalloc */
		mutex_lock(&dim_kexec_mutex);
		dim_kexec_buffer = vmalloc(kexec_buffer_size);
		if (!dim_kexec_buffer) {
			pr_err("DIM: Cannot allocate memory for restoring measurements\n");
			mutex_unlock(&dim_kexec_mutex);
			return;
		}

		dim_kexec_buffer_size = kexec_buffer_size;
		memcpy(dim_kexec_buffer, kexec_buffer, kexec_buffer_size);
		mutex_unlock(&dim_kexec_mutex);
		r = dim_free_kexec_buffer();
		if (r < 0)
			pr_warn("There is still kexec buffer not freed, r:%d", r);
		break;
	case -ENOENT:
		pr_info("DIM: No measurement list to restore, it's ok\n");
		break;
	default:
		pr_info("DIM: restoring the measurement list: %d\n", r);
		break;
	}
}

static int __init dim_init(void)
{
	/* Reserved for future extra dim init process */
	dim_load_kexec_buffer();
	return 0;
}
late_initcall(dim_init);
